﻿// Copyright (c) Microsoft. All rights reserved.

using System.Text.Json.Serialization;

namespace Microsoft.Agents.AI.Hosting.OpenAI.ChatCompletions.Models;

/// <summary>
/// Represents usage statistics for a chat completion request.
/// </summary>
internal sealed record CompletionUsage
{
    public static CompletionUsage Zero { get; } = new()
    {
        CompletionTokens = 0,
        PromptTokens = 0,
        TotalTokens = 0,
        CompletionTokensDetails = new()
        {
            AcceptedPredictionTokens = 0,
            AudioTokens = 0,
            ReasoningTokens = 0,
            RejectedPredictionTokens = 0
        },
        PromptTokensDetails = new()
        {
            AudioTokens = 0,
            CachedTokens = 0
        },
    };

    /// <summary>
    /// Number of tokens in the generated completion.
    /// </summary>
    [JsonPropertyName("completion_tokens")]
    public int? CompletionTokens { get; set; }

    /// <summary>
    /// Number of tokens in the prompt.
    /// </summary>
    [JsonPropertyName("prompt_tokens")]
    public int? PromptTokens { get; set; }

    /// <summary>
    /// Total number of tokens used in the request (prompt + completion).
    /// </summary>
    [JsonPropertyName("total_tokens")]
    public int? TotalTokens { get; set; }

    /// <summary>
    /// Breakdown of tokens used in the generated completion.
    /// </summary>
    [JsonPropertyName("completion_tokens_details")]
    public required CompletionTokensDetails CompletionTokensDetails { get; set; }

    /// <summary>
    /// Breakdown of tokens used in the prompt.
    /// </summary>
    [JsonPropertyName("prompt_tokens_details")]
    public required PromptTokensDetails PromptTokensDetails { get; set; }

    public static CompletionUsage operator +(CompletionUsage left, CompletionUsage right) => new()
    {
        CompletionTokens = left.CompletionTokens + right.CompletionTokens,
        PromptTokens = left.PromptTokens + right.PromptTokens,
        TotalTokens = left.TotalTokens + right.TotalTokens,
        CompletionTokensDetails = left.CompletionTokensDetails + right.CompletionTokensDetails,
        PromptTokensDetails = left.PromptTokensDetails + right.PromptTokensDetails
    };
}

/// <summary>
/// Breakdown of tokens used in a completion.
/// </summary>
internal sealed record CompletionTokensDetails
{
    /// <summary>
    /// When using Predicted Outputs, the number of tokens in the prediction that appeared in the completion.
    /// </summary>
    [JsonPropertyName("accepted_prediction_tokens")]
    public int AcceptedPredictionTokens { get; set; }

    /// <summary>
    /// Audio input tokens generated by the model.
    /// </summary>
    [JsonPropertyName("audio_tokens")]
    public int AudioTokens { get; set; }

    /// <summary>
    /// Tokens generated by the model for reasoning.
    /// </summary>
    [JsonPropertyName("reasoning_tokens")]
    public int ReasoningTokens { get; set; }

    /// <summary>
    /// When using Predicted Outputs, the number of tokens in the prediction that did not appear in the completion.
    /// However, like reasoning tokens, these tokens are still counted in the total completion tokens for purposes of billing,
    /// output, and context window limits.
    /// </summary>
    [JsonPropertyName("rejected_prediction_tokens")]
    public int RejectedPredictionTokens { get; set; }

    public static CompletionTokensDetails operator +(CompletionTokensDetails left, CompletionTokensDetails right) => new()
    {
        AcceptedPredictionTokens = left.AcceptedPredictionTokens + right.AcceptedPredictionTokens,
        AudioTokens = left.AudioTokens + right.AudioTokens,
        ReasoningTokens = left.ReasoningTokens + right.ReasoningTokens,
        RejectedPredictionTokens = left.RejectedPredictionTokens + right.RejectedPredictionTokens
    };
}

/// <summary>
/// Breakdown of tokens used in the prompt.
/// </summary>
internal sealed record PromptTokensDetails
{
    /// <summary>
    /// Audio input tokens present in the prompt.
    /// </summary>
    [JsonPropertyName("audio_tokens")]
    public int AudioTokens { get; set; }

    /// <summary>
    /// Cached tokens present in the prompt.
    /// </summary>
    [JsonPropertyName("cached_tokens")]
    public int CachedTokens { get; set; }

    public static PromptTokensDetails operator +(PromptTokensDetails left, PromptTokensDetails right) => new()
    {
        AudioTokens = left.AudioTokens + right.AudioTokens,
        CachedTokens = left.CachedTokens + right.CachedTokens
    };
}
